#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ChromaDB 数据诊断脚本
用于检查 memories 和 conversations 集合中的数据情况
"""

import asyncio
import sys
import os
from pathlib import Path

# 添加项目根目录到路径
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))
sys.path.insert(0, str(project_root / "src"))

# 设置环境变量
# os.chdir(project_root / "src")  # 移除以避免相对导入错误
from database.manager import DatabaseManager
from loguru import logger
from config import load_config
config = load_config()

async def diagnose_chromadb():
    """诊断ChromaDB数据情况"""
    try:
        # 初始化数据库管理器
        db_manager = DatabaseManager(config)
        await db_manager.initialize()
        
        chromadb_client = db_manager.chromadb_client
        if not chromadb_client:
            print("❌ ChromaDB客户端未初始化")
            return
            
        print("✅ ChromaDB客户端已连接")
        
        # 检查集合状态
        collections = ['memories', 'conversations', 'knowledge', 'documents']
        
        for collection_name in collections:
            print(f"\n📊 检查集合: {collection_name}")
            try:
                # 获取集合统计信息
                stats = await chromadb_client.get_collection_stats(collection_name)
                print(f"  - 文档数量: {stats.get('count', 0)}")
                
                # 获取一些示例数据
                if stats.get('count', 0) > 0:
                    sample_data = await chromadb_client.get_documents(
                        collection_name, limit=3
                    )
                    
                    print(f"  - 示例文档数: {len(sample_data.get('documents', []))}")
                    
                    # 显示元数据示例
                    metadatas = sample_data.get('metadatas', [])
                    if metadatas:
                        print(f"  - 示例元数据: {metadatas[0]}")
                        
                        # 检查user_id字段
                        user_ids = [meta.get('user_id') for meta in metadatas if meta.get('user_id')]
                        if user_ids:
                            print(f"  - 发现用户ID: {set(user_ids)}")
                        else:
                            print(f"  - ⚠️  未发现user_id字段")
                else:
                    print(f"  - ⚠️  集合为空")
                    
            except Exception as e:
                print(f"  - ❌ 检查失败: {e}")
        
        # 测试搜索功能
        print(f"\n🔍 测试搜索功能")
        
        # 测试memory搜索
        try:
            memory_results = await chromadb_client.search_memories(
                user_id="test_user", 
                query="测试", 
                n_results=5
            )
            print(f"  - memory搜索结果: {len(memory_results)} 条")
            if memory_results:
                print(f"    示例: {memory_results[0].get('content', '')[:50]}...")
        except Exception as e:
            print(f"  - ❌ memory搜索失败: {e}")
            
        # 测试conversation搜索
        try:
            conv_results = await chromadb_client.search_conversations(
                query="测试", 
                user_id="test_user", 
                limit=5
            )
            print(f"  - conversation搜索结果: {len(conv_results)} 条")
            if conv_results:
                print(f"    示例: {conv_results[0].get('content', '')[:50]}...")
        except Exception as e:
            print(f"  - ❌ conversation搜索失败: {e}")
            
        # 测试knowledge搜索（对比）
        try:
            knowledge_results = await chromadb_client.search_documents(
                'knowledge', 
                "测试", 
                n_results=5
            )
            docs = knowledge_results.get('documents', [])
            print(f"  - knowledge搜索结果: {len(docs)} 条")
            if docs:
                print(f"    示例: {docs[0][:50]}...")
        except Exception as e:
            print(f"  - ❌ knowledge搜索失败: {e}")
            
        # 测试无过滤条件的搜索
        print(f"\n🔍 测试无过滤条件搜索")
        try:
            # memories集合无过滤条件搜索
            raw_memory_results = await chromadb_client.search_documents(
                'memories', 
                "测试", 
                n_results=5
            )
            docs = raw_memory_results.get('documents', [])
            print(f"  - memories原始搜索结果: {len(docs)} 条")
            
            # conversations集合无过滤条件搜索
            raw_conv_results = await chromadb_client.search_documents(
                'conversations', 
                "测试", 
                n_results=5
            )
            docs = raw_conv_results.get('documents', [])
            print(f"  - conversations原始搜索结果: {len(docs)} 条")
            
        except Exception as e:
            print(f"  - ❌ 原始搜索失败: {e}")
            
    except Exception as e:
        print(f"❌ 诊断失败: {e}")
        import traceback
        traceback.print_exc()
    finally:
        if 'db_manager' in locals():
            await db_manager.close()

if __name__ == "__main__":
    asyncio.run(diagnose_chromadb())